require "Q2DViewer"
require "luaQt"

local nl = require"luaNumeric"
--local GLUI = require"luaGLShapes"
local nl = require"luaNumeric"
-- local DataSheet = require("DataSheet")


----------------------------------------------------------------------------------------------------------------------------------------
--初始化参数
local FEM2D={
	freq = 1e7;
	beta0 = 1;
	alpha0 = 1;
};
----------------------------------------------------------------------------------------------------------------------------------------

----------------------------------------------------------------------------------------------------------------------------------------
--创建表和创建元表
function FEM2D:new (o)
	o = o or {};
	self.__index = self;
	setmetatable(o,self)
	o:init();
	return o;
end
----------------------------------------------------------------------------------------------------------------------------------------


----------------------------------------------------------------------------------------------------------------------------------------
--初始化
function FEM2D:init ()
	self.omega = 2*pi*self.freq;
	self.k0=self.omega/c0;
	print("k0="..self.k0)
	self.mesh = Geometry.create_mesh_soup();
	self.mesh:is_outline(true)
end
----------------------------------------------------------------------------------------------------------------------------------------


----------------------------------------------------------------------------------------------------------------------------------------
--设置形状，网格剖分
function FEM2D:set_shapes (shapes)
	self.mesh = Geometry.create_mesh_soup();
	self.mesh:is_outline(true)
	
	local mesh = self.mesh;
	mesh:shapes(shapes);
	mesh:create_mesh();
	
	------------------------------------------------------------
	--所有节点坐标，和节点标号
	self.pts = mesh:points();
	self.npts = #self.pts
	------------------------------------------------------------
	
	------------------------------------------------------------
	--所有线段信息和线段编组信息
	self.segs_infos = mesh:segments_info();
	print('Segs info size:', #self.segs_infos);
	self.segs_groups_tags = mesh:segments_groups_tags();
	self.segments_groups_nodes = mesh:segments_groups_nodes();
	self.nedges = #self.segs_infos;
	------------------------------------------------------------
	
	------------------------------------------------------------
	--所有三角形信息
	self.trigs_infos = mesh:triangles_info();
	self.trigs_groups_tag = mesh:triangles_groups_tags();
	self.ntrigs = #self.trigs_infos;
	print( 'trigs infos size :',#self.trigs_infos);
	print( 'trigs groups size :',#self.trigs_groups_tag);
	------------------------------------------------------------
	
	------------------------------------------------------------
	--清空所有边界条件信息
	self.N_bc = {};
	self.D_bc = {};
	------------------------------------------------------------
end
----------------------------------------------------------------------------------------------------------------------------------------

----------------------------------------------------------------------------------------------------------------------------------------
--初始化材料参数
function FEM2D:init_material()
	local ntrigs =self.ntrigs;
	for kk=1,ntrigs do
		local info = self.trigs_infos[kk];
		info:alpha(self.alpha0)
		info:beta(self.beta0)
	end
	
	local nedges =self.nedges;
	for kk=1,nedges do
		local info = self.segs_infos[kk];
		info:gamma(0)
		info:q(0)
	end
end
----------------------------------------------------------------------------------------------------------------------------------------

----------------------------------------------------------------------------------------------------------------------------------------
--设置材料参数
function FEM2D:set_epsr_mur(tag,epsr,mur)
	local trigs_group = self.trigs_groups_tag[tag];
	assert(tag > 0 and tag <= #self.trigs_groups_tag, "out of range:"..tag..' not in range [1,'.. #self.trigs_groups_tag .. ']')
	for kk=1,#trigs_group do
		local info = self.trigs_infos[trigs_group [kk]];
		info:alpha(1/mur)
		info:beta(epsr)
	end
end
----------------------------------------------------------------------------------------------------------------------------------------

----------------------------------------------------------------------------------------------------------------------------------------
--初始化矩阵
function FEM2D:init_matrix()
	local npt= self.npts;
	
	self.spF =  nl.new_sparse_cmatrix(npt,npt)
	self.vecb =  nl.new_eigen_vectorxcd(npt)
	self.vecx =  nl.new_eigen_vectorxcd(npt)
	self.vecd =  nl.new_eigen_vectorxcd(npt)
	
	--self.spF =  nl.new_sparse_matrix(npt,npt)
	--self.vecb =  nl.new_eigen_vectorxd(npt)
	--self.vecx =  nl.new_eigen_vectorxd(npt)
	--self.vecd =  nl.new_eigen_vectorxd(npt)
	
	--self.spF =  nl.new_sparse_matrix(npt,npt)
	--self.vecb =  nl.new_eigen_vectorxd(npt)
	--self.vecx =  nl.new_eigen_vectorxd(npt)

	--self.spF:set_identity();
	self.spF:set_zero();
	self.vecb:set_zero();
	self.vecx:set_zero();
	self.vecd:set_zero();
	--print(self.spF)
end
----------------------------------------------------------------------------------------------------------------------------------------


----------------------------------------------------------------------------------------------------------------------------------------
--设置矩阵
function FEM2D:set_matrix()
	
	function fij(v1,v2,b,c)
		return b[v1]*b[v2]+c[v1]*c[v2];
	end
	local spF = self.spF
	local vecb = self.vecb

	
	local ntrigs = self.ntrigs;
	local trigs_infos = self.trigs_infos;

	for kk=1,ntrigs do
		local info = trigs_infos[kk];
		local pt_tag = info:nodes_tag();
		local area = info:area();
		local a = info:a();
		local b = info:b();
		local c = info:c();
		local fe = info:f():real();
		local alpha = info:alpha():real();
		local beta = info:beta():real();
		for ii=1,3 do		
			for jj=1,3 do
				spF[{pt_tag[ii],pt_tag[jj]}] = spF[{pt_tag[ii],pt_tag[jj]}]  +
					1.0/(4.0*area)*( alpha*b[ii]*b[jj] + alpha*c[ii]*c[jj] )
					+area/12*beta*(1+nl. kronecker_delta(ii,jj));
			end
			vecb[pt_tag[ii]] = vecb[pt_tag[ii]] + area/3.0*fe;
		end
	end
end
----------------------------------------------------------------------------------------------------------------------------------------


----------------------------------------------------------------------------------------------------------------------------------------
--强加第一类边界条件：1.记录第一类边界条件
function FEM2D:set_D_BCs(group_tag, phi)
	local bcs2 = self.segments_groups_nodes[group_tag];
	for kk=1,#bcs2 do
		tag = bcs2[kk];
		self.D_bc[tag] = phi;		
	end
end
----------------------------------------------------------------------------------------------------------------------------------------

----------------------------------------------------------------------------------------------------------------------------------------
--强加第一类边界条件：2.设置矩阵
function FEM2D:set_D_Matrix()
	local npts = self.npts
	
	for ii=1,npts do		
		for jj,v in pairs(self.D_bc) do
			self.vecb[ii] = self.vecb[ii] - self.spF[{ii,jj}] * v;
		end		
	end
		
	for ii,v in pairs(self.D_bc) do
		self.vecb[ii] = v ;
		self.spF:set_row(ii,0);
		self.spF:set_col(ii,0);
		self.spF[{ii,ii}] = 1;
	end

end
----------------------------------------------------------------------------------------------------------------------------------------


----------------------------------------------------------------------------------------------------------------------------------------
--设置第三类边界条件
function FEM2D:set_N_BCs(group_tag, q,gamma)
	local bcs2 = self.segs_groups_tags[group_tag];
	for kk=1,#bcs2 do
		local tag = bcs2[kk];
		self.N_bc[tag] = {q,gamma};		
	end
end
----------------------------------------------------------------------------------------------------------------------------------------

----------------------------------------------------------------------------------------------------------------------------------------
--第三类边界条件
function FEM2D:set_N_Matrix()
	--print(self.segs_infos)
	local segs_infos = self.segs_infos;
	for kk,vv in pairs(self.N_bc) do
		local info = segs_infos[kk];
		local pt_tag = info:nodes_tag();
		local ls = info:length();
		local q = vv[1];
		local r = vv[2];
		for ii=1,2 do		
			for jj=1,2 do
				self.spF[{pt_tag[ii],pt_tag[jj]}] = self.spF[{pt_tag[ii],pt_tag[jj]}]  +
					r*ls/(6.0)*(1+ nl. kronecker_delta(ii,jj));
			end
			self.vecb[pt_tag[ii]] =  self.vecb[pt_tag[ii]] +
				q*ls/2.0
		end
	end
end
----------------------------------------------------------------------------------------------------------------------------------------

----------------------------------------------------------------------------------------------------------------------------------------
--求解
function FEM2D:solve()
	self:set_N_Matrix();
	self:set_D_Matrix();
	local ts = os.clock()	
	assert( self.spF:solve_lu(self.vecb,self.vecx)>=0 , "solve failed")
	--assert( self.spF:solve_ldlt(self.vecb,self.vecx)>=0 , "solve failed")
	print(string.format("elapsed time: %.2f\n", os.clock() - ts))
end
----------------------------------------------------------------------------------------------------------------------------------------

----------------------------------------------------------------------------------------------------------------------------------------
--后处理
function FEM2D:post_calculate()
	local npts = self.npts;
	local fields = self.mesh:points_vals();
	fields:resize( npts )	
	for ii=1,npts do
		fields[ii] = nl.real(self.vecx[ii]);
	end
end
----------------------------------------------------------------------------------------------------------------------------------------

return FEM2D;